import cdf.jce.provider.CDFJCE;
import java.nio.charset.StandardCharsets;
import java.security.cert.CertificateFactory;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.security.cert.X509Certificate;
import java.security.cert.CertificateException;
import java.security.*;
import java.security.cert.*;
import java.util.*;
import java.io.*;

public class JCECrypto {
	// 加密机提供者名称
	private static final String PROVIDER = "CDFJCE";
	private static final String SM1_ALGORITHM = "SM1";
	private static final String SM4_ALGORITHM = "SM4";

	// 密钥长度常量
	private static final int KEY_LENGTH = 128;
	private static final int KEY_BYTES = KEY_LENGTH / 8;

	// 初始化加密机提供者
	static {
		try {
			if (Security.getProvider(PROVIDER) == null) {
				Security.addProvider(new cdf.jce.provider.CDFJCE());
			}
		} catch (Exception e) {
			System.err.println("加密机提供者初始化失败: " + e.getMessage());
		}
	}

	// ====================== 证书相关接口 ======================

	/**
	 * 使用证书进行SM2加密
	 * 
	 * @param base64Certificate Base64编码的证书数据
	 * @param plaintext         明文数据
	 * @return 加密后的密文
	 */
	public static byte[] encryptWithCertificate(String base64Certificate, byte[] plaintext) {
		try {
			// 1. 验证参数
			if (base64Certificate == null || base64Certificate.isEmpty()) {
				throw new IllegalArgumentException("证书不能为空");
			}

			if (plaintext == null || plaintext.length == 0) {
				throw new IllegalArgumentException("明文不能为空");
			}

			// 2. 解码Base64证书数据
			byte[] certificateBytes = Base64.getDecoder().decode(base64Certificate);

			// 3. 获取证书工厂实例
			CertificateFactory certificateFactory = CertificateFactory.getInstance("X509", PROVIDER);

			// 4. 生成证书对象
			X509Certificate certificate = (X509Certificate) certificateFactory.generateCertificate(
					new ByteArrayInputStream(certificateBytes));

			// 5. 使用公钥加密
			return sm2Encrypt(certificate.getPublicKey(), plaintext);
		} catch (CertificateException e) {
			throw new CryptoRuntimeException("证书处理错误", e);
		} catch (NoSuchProviderException e) {
			throw new CryptoRuntimeException("找不到安全提供者", e);
		} catch (IllegalArgumentException e) {
			throw new CryptoRuntimeException("Base64解码错误", e);
		} catch (Exception e) {
			throw new CryptoRuntimeException("加密失败", e);
		}
	}

	/**
	 * SM2公钥加密
	 * 
	 * @param publicKey 公钥
	 * @param plaintext 明文数据
	 * @return 加密后的密文
	 */
	public static byte[] sm2Encrypt(PublicKey publicKey, byte[] plaintext) {
		try {
			if (publicKey == null) {
				throw new IllegalArgumentException("公钥不能为空");
			}

			if (plaintext == null || plaintext.length == 0) {
				throw new IllegalArgumentException("明文不能为空");
			}

			CDFSM2 sm2 = new CDFSM2();
			return sm2.ExternalSM2Enc(publicKey, plaintext);
		} catch (Exception e) {
			throw new CryptoRuntimeException("SM2加密失败", e);
		}
	}

	/**
	 * 符合国密标准的密钥派生 (GB/T 32918.4-2016)
	 * 
	 * @param masterKey             主密钥 (16字节)
	 * @param diversificationFactor 分散因子
	 * @param keyLength             派生密钥长度
	 */
	public static byte[] kdf(byte[] masterKey, byte[] diversificationFactor, int keyLength) {
		try {
			MessageDigest md = MessageDigest.getInstance("SM3", PROVIDER);
			md.update(masterKey);
			md.update(diversificationFactor);
			byte[] derived = md.digest();
			return Arrays.copyOf(derived, keyLength / 8);
		} catch (Exception e) {
			throw new CryptoRuntimeException("KDF执行失败", e);
		}
	}

	/**
	 * 派生SM4密钥
	 */
	public static SecretKey deriveSM4Key(byte[] masterKey, byte[] diversificationFactor) {
		// 参数校验
		if (masterKey == null || masterKey.length != KEY_BYTES) {
			throw new IllegalArgumentException("主密钥必须是" + KEY_LENGTH + "位(" + KEY_BYTES + "字节)");
		}

		try {
			// TODO: 启用真正的KDF派生
			byte[] derivedKey = kdf(masterKey, diversificationFactor, KEY_LENGTH);
			return new SecretKeySpec(derivedKey, SM4_ALGORITHM);
		} finally {
			// 清除敏感数据
			if (diversificationFactor != null) {
				Arrays.fill(diversificationFactor, (byte) 0);
			}
		}
	}

	/**
	 * 生成SM1密钥
	 */
	public static SecretKey generateSM1Key() {
		try {
			// 1. 获取安全随机数生成器
			SecureRandom ran = SecureRandom.getInstance("TrueRandom", PROVIDER);

			// 2. 创建SM1密钥生成器
			KeyGenerator skg = KeyGenerator.getInstance(SM1_ALGORITHM, PROVIDER);

			// 3. 初始化密钥生成器
			skg.init(KEY_LENGTH, ran);

			// 4. 生成SM1密钥
			SecretKey key = skg.generateKey();

			// 5. 验证密钥合规性
			validateKey(key, SM1_ALGORITHM);

			return key;
		} catch (NoSuchAlgorithmException | NoSuchProviderException e) {
			throw new CryptoRuntimeException("SM1密钥生成失败：算法或提供者不可用", e);
		} catch (Exception e) {
			throw new CryptoRuntimeException("SM1密钥生成发生未预期错误", e);
		}
	}

	/**
	 * 生成SM4密钥
	 */
	public static SecretKey generateSM4Key() {
		try {
			// 1. 获取安全随机数生成器
			SecureRandom ran = SecureRandom.getInstance("TrueRandom", PROVIDER);

			// 2. 创建SM4密钥生成器
			KeyGenerator skg = KeyGenerator.getInstance(SM4_ALGORITHM, PROVIDER);

			// 3. 初始化密钥生成器
			skg.init(KEY_LENGTH, ran);

			// 4. 生成SM4密钥
			SecretKey key = skg.generateKey();

			// 5. 验证密钥合规性
			validateKey(key, SM4_ALGORITHM);

			return key;
		} catch (NoSuchAlgorithmException | NoSuchProviderException e) {
			throw new CryptoRuntimeException("SM4密钥生成失败：算法或提供者不可用", e);
		} catch (Exception e) {
			throw new CryptoRuntimeException("SM4密钥生成发生未预期错误", e);
		}
	}

	/**
	 * 验证密钥合规性
	 */
	private static void validateKey(SecretKey key, String expectedAlgorithm)
			throws InvalidKeyException { // 添加异常声明

		if (key == null) {
			throw new InvalidKeyException("生成密钥失败：返回null");
		}

		byte[] keyBytes = key.getEncoded();
		final int KEY_BYTES = 16; // SM4 固定 16 字节密钥

		if (keyBytes == null || keyBytes.length != KEY_BYTES) {
			throw new InvalidKeyException("无效密钥长度: " +
					(keyBytes != null ? keyBytes.length * 8 : "null") + "位，应为128位");
		}

		if (!expectedAlgorithm.equals(key.getAlgorithm())) {
			throw new InvalidKeyException("密钥算法不匹配: 期望" +
					expectedAlgorithm + ", 实际" + key.getAlgorithm());
		}
	}

	// ====================== 测试方法 ======================

	public static void main(String[] args) {
		if (args.length < 2) {
			printHelp();
			System.exit(1);
		}

		String command = args[0];
		switch (command) {
			case "encrypt":
				handleEncryptCommand(args);
				break;
			case "genkey":
				handleGenKeyCommand(args);
				break;
			case "version":
				printVersion();
				break;
			default:
				printHelp();
				System.exit(1);
		}
	}

	private static void printHelp() {
		System.out.println("JCECrypto - Java 加密服务 CLI");
		System.out.println("用法:");
		System.out.println("  encrypt -cert <Base64证书> -data <Base64明文>");
		System.out.println("  genkey -type <SM1|SM4|DERIVED_SM4> [-diver <Base64分散因子>] [-rootkey <Base64根密钥>]");
		System.out.println("  version");
		System.out.println("示例:");
		System.out.println("  java -jar JCECrypto.jar encrypt -cert MIIC.. -data VEVTVCBEQVRB");
		System.out.println("  java -jar JCECrypto.jar genkey -type DERIVED_SM4 -rootkey MDEyMz.. -diver RkM..");
	}

	private static void printVersion() {
		try {
			Package pkg = JCECrypto.class.getPackage();
			String version = pkg != null ? pkg.getImplementationVersion() : "1.0";
			System.out.println("JCECrypto v" + (version != null ? version : "1.0"));
			System.out.println("使用加密机 SDK: CDFJCE 2.1.2");
		} catch (Exception e) {
			System.out.println("JCECrypto v1.0");
		}
	}

	private static void handleEncryptCommand(String[] args) {
		String certificate = null;
		String data = null;

		for (int i = 1; i < args.length; i++) {
			if ("-cert".equals(args[i]) && i + 1 < args.length) {
				certificate = args[++i];
			} else if ("-data".equals(args[i]) && i + 1 < args.length) {
				data = args[++i];
			}
		}

		try {
			// 参数检查
			if (certificate == null || certificate.isEmpty()) {
				throw new IllegalArgumentException("缺少证书参数");
			}

			if (data == null || data.isEmpty()) {
				throw new IllegalArgumentException("缺少数据参数");
			}

			// Base64解码
			byte[] plaintext = Base64.getDecoder().decode(data);
			byte[] ciphertext = encryptWithCertificate(certificate, plaintext);

			if (ciphertext == null || ciphertext.length == 0) {
				throw new RuntimeException("加密失败：返回空结果");
			}

			// 输出加密结果
			System.out.println("加密成功(hex)：");
			ComFun.printfHexString(ciphertext);
			System.out.println("加密成功(base64)：");
			System.out.println(Base64.getEncoder().encodeToString(ciphertext));
			System.exit(0);
		} catch (IllegalArgumentException e) {
			System.err.println("参数错误: " + e.getMessage());
			System.exit(1);
		} catch (Exception e) {
			System.err.println("加密过程中出错: " + e.getMessage());
			e.printStackTrace();
			System.exit(1);
		}
	}

	private static void handleGenKeyCommand(String[] args) {
		String type = null;
		String diver = null;
		String rootkey = null;

		for (int i = 1; i < args.length; i++) {
			if ("-type".equals(args[i]) && i + 1 < args.length) {
				type = args[++i];
			} else if ("-diver".equals(args[i]) && i + 1 < args.length) {
				diver = args[++i];
			} else if ("-rootkey".equals(args[i]) && i + 1 < args.length) {
				rootkey = args[++i];
			}
		}

		try {
			// 参数检查
			if (type == null || type.isEmpty()) {
				throw new IllegalArgumentException("缺少密钥类型参数");
			}

			SecretKey key = null;
			byte[] keyBytes = null;
			byte[] diverBytes = null;

			try {
				switch (type.toUpperCase()) {
					case "SM1":
						key = generateSM1Key();
						break;
					case "SM4":
						key = generateSM4Key();
						break;
					case "DERIVED_SM4":
						if (rootkey == null || rootkey.isEmpty()) {
							throw new IllegalArgumentException("需要-rootkey参数");
						}

						if (diver == null || diver.isEmpty()) {
							throw new IllegalArgumentException("需要-diver参数");
						}

						keyBytes = Base64.getDecoder().decode(rootkey);
						diverBytes = Base64.getDecoder().decode(diver);

						// 验证长度
						if (keyBytes.length != KEY_BYTES) {
							throw new IllegalArgumentException("根密钥必须是" + KEY_LENGTH + "位(" + KEY_BYTES + "字节)");
						}

						key = deriveSM4Key(keyBytes, diverBytes);
						break;
					default:
						throw new IllegalArgumentException("不支持的密钥类型: " + type);
				}

				if (key == null) {
					throw new RuntimeException("密钥生成失败: 返回null");
				}

				byte[] encodedKey = key.getEncoded();

				if (encodedKey == null || encodedKey.length == 0) {
					throw new RuntimeException("密钥生成失败: 无法获取密钥内容");
				}

				// 输出密钥
				System.out.println("密钥生成成功(hex)：");
				ComFun.printfHexString(encodedKey);
				System.out.println("密钥生成成功(base64)：");
				System.out.println(Base64.getEncoder().encodeToString(encodedKey));
				System.exit(0);
			} finally {
				// 清除敏感数据
				if (keyBytes != null) {
					Arrays.fill(keyBytes, (byte) 0);
				}
				if (diverBytes != null) {
					Arrays.fill(diverBytes, (byte) 0);
				}
			}
		} catch (IllegalArgumentException e) {
			System.err.println("参数错误: " + e.getMessage());
			System.exit(1);
		} catch (Exception e) {
			System.err.println("密钥生成过程中出错: " + e.getMessage());
			e.printStackTrace();
			System.exit(1);
		}
	}

	/**
	 * 自定义加密运行时异常
	 */
	private static class CryptoRuntimeException extends RuntimeException {
		public CryptoRuntimeException(String message) {
			super(message);
		}

		public CryptoRuntimeException(String message, Throwable cause) {
			super(message, cause);
		}
	}
}